#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

#define MAX_PACKETS 1000
#define STATS_INTERVAL 60  // 统计间隔(秒)

// 统计结构体
typedef struct {
    char mac[18];
    char ip[16];
    long long total_bytes_sent;
    long long total_bytes_received;
    time_t last_update;
} host_stats_t;

host_stats_t stats[1000];
int stats_count = 0;

// 写入基础日志
void write_basic_log(FILE *file, const char *timestamp, const char *src_mac, 
                    const char *src_ip, const char *dst_mac, 
                    const char *dst_ip, int frame_len) {
    fprintf(file, "%s,%s,%s,%s,%s,%d\n", 
            timestamp, src_mac, src_ip, dst_mac, dst_ip, frame_len);
    fflush(file);
}

// 查找或创建统计记录
host_stats_t* find_or_create_stats(const char *mac, const char *ip) {
    for (int i = 0; i < stats_count; i++) {
        if (strcmp(stats[i].mac, mac) == 0 && strcmp(stats[i].ip, ip) == 0) {
            return &stats[i];
        }
    }
    
    if (stats_count < 1000) {
        host_stats_t *new_stat = &stats[stats_count++];
        strcpy(new_stat->mac, mac);
        strcpy(new_stat->ip, ip);
        new_stat->total_bytes_sent = 0;
        new_stat->total_bytes_received = 0;
        new_stat->last_update = time(NULL);
        return new_stat;
    }
    return NULL;
}

// 更新统计数据
void update_stats(const char *src_mac, const char *src_ip, 
                 const char *dst_mac, const char *dst_ip, int frame_len) {
    host_stats_t *src_stat = find_or_create_stats(src_mac, src_ip);
    host_stats_t *dst_stat = find_or_create_stats(dst_mac, dst_ip);
    
    if (src_stat) {
        src_stat->total_bytes_sent += frame_len;
        src_stat->last_update = time(NULL);
    }
    if (dst_stat) {
        dst_stat->total_bytes_received += frame_len;
        dst_stat->last_update = time(NULL);
    }
}

// 输出统计信息
void print_statistics(FILE *stats_file) {
    time_t now = time(NULL);
    char timestamp[20];
    strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", localtime(&now));
    
    fprintf(stats_file, "\n=== 流量统计 %s ===\n", timestamp);
    
    fprintf(stats_file, "\n发送流量统计:\n");
    fprintf(stats_file, "%-18s %-15s %-12s\n", "MAC地址", "IP地址", "发送字节数");
    for (int i = 0; i < stats_count; i++) {
        if (stats[i].total_bytes_sent > 0) {
            fprintf(stats_file, "%-18s %-15s %-12lld\n", 
                    stats[i].mac, stats[i].ip, stats[i].total_bytes_sent);
        }
    }
    
    fprintf(stats_file, "\n接收流量统计:\n");
    fprintf(stats_file, "%-18s %-15s %-12s\n", "MAC地址", "IP地址", "接收字节数");
    for (int i = 0; i < stats_count; i++) {
        if (stats[i].total_bytes_received > 0) {
            fprintf(stats_file, "%-18s %-15s %-12lld\n", 
                    stats[i].mac, stats[i].ip, stats[i].total_bytes_received);
        }
    }
    fprintf(stats_file, "\n");
    fflush(stats_file);
}

// 数据包处理回调函数
void packet_handler(u_char *user_data, const struct pcap_pkthdr *pkthdr, const u_char *packet) {
    static time_t last_stats_time = 0;
    time_t current_time = time(NULL);
    
    // 获取时间戳
    char timestamp[20];
    struct tm *lt = localtime(&pkthdr->ts.tv_sec);
    strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", lt);
    
    // 解析以太网帧头
    struct ethhdr *eth_header = (struct ethhdr*)packet;
    
    // 获取MAC地址
    char src_mac[18], dst_mac[18];
    snprintf(src_mac, sizeof(src_mac), "%02X-%02X-%02X-%02X-%02X-%02X",
             eth_header->h_source[0], eth_header->h_source[1],
             eth_header->h_source[2], eth_header->h_source[3],
             eth_header->h_source[4], eth_header->h_source[5]);
             
    snprintf(dst_mac, sizeof(dst_mac), "%02X-%02X-%02X-%02X-%02X-%02X",
             eth_header->h_dest[0], eth_header->h_dest[1],
             eth_header->h_dest[2], eth_header->h_dest[3],
             eth_header->h_dest[4], eth_header->h_dest[5]);
    
    // 检查是否为IP数据包
    if (ntohs(eth_header->h_proto) == ETH_P_IP) {
        // 解析IP头
        struct iphdr *ip_header = (struct iphdr*)(packet + sizeof(struct ethhdr));
        
        char src_ip[16], dst_ip[16];
        strcpy(src_ip, inet_ntoa(*(struct in_addr*)&ip_header->saddr));
        strcpy(dst_ip, inet_ntoa(*(struct in_addr*)&ip_header->daddr));
        
        // 写入基础日志
        FILE *log_file = (FILE*)user_data;
        write_basic_log(log_file, timestamp, src_mac, src_ip, dst_mac, dst_ip, pkthdr->len);
        
        // 更新统计数据
        update_stats(src_mac, src_ip, dst_mac, dst_ip, pkthdr->len);
    } else {
        // 非IP包，只记录MAC信息
        FILE *log_file = (FILE*)user_data;
        fprintf(log_file, "%s,%s,-,%s,-,%d\n", 
                timestamp, src_mac, dst_mac, pkthdr->len);
        fflush(log_file);
    }
    
    // 定期输出统计信息
    if (current_time - last_stats_time >= STATS_INTERVAL) {
        FILE *stats_file = fopen("network_stats.txt", "a");
        if (stats_file) {
            print_statistics(stats_file);
            fclose(stats_file);
        }
        last_stats_time = current_time;
    }
}

// 获取网络设备列表并让用户选择
char* select_network_device() {
    pcap_if_t *alldevs;
    pcap_if_t *d;
    char errbuf[PCAP_ERRBUF_SIZE];
    int i = 0;
    int selection;
    static char selected_device[256];
    
    // 获取设备列表
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        fprintf(stderr, "获取网络设备列表失败: %s\n", errbuf);
        return NULL;
    }
    
    // 显示设备列表
    printf("可用的网络设备:\n");
    for (d = alldevs; d != NULL; d = d->next) {
        printf("%d. %s", ++i, d->name);
        if (d->description)
            printf(" (%s)\n", d->description);
        else
            printf(" (无描述)\n");
    }
    
    if (i == 0) {
        printf("未找到网络设备! 请检查权限。\n");
        pcap_freealldevs(alldevs);
        return NULL;
    }
    
    // 让用户选择设备
    printf("\n请选择要监听的设备 (1-%d): ", i);
    scanf("%d", &selection);
    
    if (selection < 1 || selection > i) {
        printf("无效的选择!\n");
        pcap_freealldevs(alldevs);
        return NULL;
    }
    
    // 找到选择的设备
    for (d = alldevs, i = 0; d != NULL; d = d->next) {
        if (++i == selection) {
            strcpy(selected_device, d->name);
            break;
        }
    }
    
    pcap_freealldevs(alldevs);
    return selected_device;
}

int main() {
    pcap_t *handle;
    char errbuf[PCAP_ERRBUF_SIZE];
    char *device;
    
    printf("=== 网络流量侦听器 ===\n");
    
    // 获取网络设备
    device = select_network_device();
    if (device == NULL) {
        fprintf(stderr, "无法选择网络设备\n");
        return 1;
    }
    printf("使用网络设备: %s\n", device);
    
    // 打开网络设备进行侦听
    handle = pcap_open_live(device, BUFSIZ, 1, 1000, errbuf);
    if (handle == NULL) {
        fprintf(stderr, "无法打开设备 %s: %s\n", device, errbuf);
        return 1;
    }
    
    // 打开日志文件
    FILE *log_file = fopen("network_traffic.csv", "w");
    if (!log_file) {
        fprintf(stderr, "无法创建日志文件\n");
        pcap_close(handle);
        return 1;
    }
    
    // 写入CSV头部
    fprintf(log_file, "时间,源MAC,源IP,目标MAC,目标IP,帧长度\n");
    fflush(log_file);
    
    // 创建统计文件
    FILE *stats_file = fopen("network_stats.txt", "w");
    if (stats_file) {
        fprintf(stats_file, "网络流量统计日志\n");
        fprintf(stats_file, "开始时间: %s\n", __DATE__ " " __TIME__);
        fclose(stats_file);
    }
    
    printf("开始侦听网络流量...\n");
    printf("日志文件: network_traffic.csv\n");
    printf("统计文件: network_stats.txt\n");
    printf("按Ctrl+C停止侦听\n");
    
    // 开始捕获数据包
    pcap_loop(handle, -1, packet_handler, (u_char*)log_file);
    
    // 清理资源
    pcap_close(handle);
    fclose(log_file);
    
    printf("侦听结束\n");
    return 0;
}
